import cv2
import numpy as np
import argparse
from ultralytics import YOLO

# 全局变量
selected_box = None
tracking = False
target_class = None
click_x, click_y = -1, -1
mouse_clicked = False
debug = False


# 鼠标回调函数
def mouse_callback(event, x, y, flags, param):
    global click_x, click_y, mouse_clicked, tracking
    if event == cv2.EVENT_LBUTTONDOWN:
        click_x, click_y = x, y
        mouse_clicked = True
        if debug:
            print(f"Left click at: ({x}, {y})")
    elif event == cv2.EVENT_RBUTTONDOWN:
        tracking = False
        if debug:
            print("Tracking stopped")


# 解析参数
def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--video", type=str, default="26075922957-1-192.mp4", help="Video source (0 for webcam)")
    parser.add_argument("--confidence", type=float, default=0.4, help="Confidence threshold")
    parser.add_argument("--model", type=str, default="yolov10n.pt", help="YOLO model path")
    parser.add_argument("--tracker", type=str, default="kcf",
                        choices=["mosse", "csrt", "kcf"], help="Tracker type")
    parser.add_argument("--debug", action="store_true", help="Enable debug mode")
    return parser.parse_args()



def create_tracker(tracker_type):
    """使用cv2.legacy模块创建跟踪器，兼容新版OpenCV"""
    try:
        if tracker_type == "mosse":
            return cv2.legacy.TrackerMOSSE_create()  # 注意这里的legacy前缀
        elif tracker_type == "csrt":
            return cv2.legacy.TrackerCSRT_create()
        elif tracker_type == "kcf":
            return cv2.legacy.TrackerKCF_create()
        else:
            print(f"Unsupported tracker: {tracker_type}, using MOSSE")
            return cv2.legacy.TrackerMOSSE_create()
    except AttributeError as e:
        print(f"Failed to create tracker: {e}")
        print("Please check OpenCV installation (must include opencv-contrib-python)")
        return None


# 检测目标
def detect_objects(frame, model, confidence):
    results = model(frame, conf=confidence)
    boxes = []  # (x, y, w, h)
    class_names = []
    for result in results:
        for box in result.boxes:
            x1, y1, x2, y2 = box.xyxy[0].tolist()
            boxes.append((int(x1), int(y1), int(x2 - x1), int(y2 - y1)))
            class_names.append(model.names[int(box.cls[0])])
    return boxes, class_names


# 检查点击是否在目标框内
def point_in_box(x, y, box):
    bx, by, bw, bh = box
    return bx <= x <= bx + bw and by <= y <= by + bh


# 主函数
def main():
    global selected_box, tracking, target_class, mouse_clicked, debug

    args = parse_args()
    debug = args.debug

    # 验证跟踪器是否可用
    test_tracker = create_tracker(args.tracker)
    if test_tracker is None:
        return  # 跟踪器无法创建，直接退出
    del test_tracker  # 释放测试跟踪器

    # 加载YOLO模型
    try:
        model = YOLO(args.model)
    except Exception as e:
        print(f"Failed to load YOLO model: {e}")
        return

    # 打开视频源
    cap = cv2.VideoCapture(args.video)
    if not cap.isOpened():
        print(f"Cannot open video source: {args.video}")
        return

    # 创建单个窗口
    window_name = "Tracker"
    cv2.namedWindow(window_name)
    cv2.setMouseCallback(window_name, mouse_callback)

    tracker = None
    font = cv2.FONT_HERSHEY_SIMPLEX

    while True:
        ret, frame = cap.read()
        if not ret:
            print("End of video")
            break

        # 检测目标
        boxes, class_names = detect_objects(frame, model, args.confidence)

        # 处理点击选择目标
        if mouse_clicked and not tracking:
            for i, (box, cls) in enumerate(zip(boxes, class_names)):
                if point_in_box(click_x, click_y, box):
                    selected_box = box
                    target_class = cls
                    tracker = create_tracker(args.tracker)
                    if tracker:
                        tracker.init(frame, selected_box)
                        tracking = True
                        print(f"Tracking {target_class} with {args.tracker}")
                    break
            mouse_clicked = False

        # 跟踪逻辑
        if tracking and tracker:
            ok, bbox = tracker.update(frame)
            if ok:
                x, y, w, h = [int(v) for v in bbox]
                cv2.rectangle(frame, (x, y), (x + w, y + h), (0, 255, 0), 2)
                cv2.putText(frame, f"Track: {target_class}", (x, y - 10),
                            font, 0.5, (0, 255, 0), 2)
            else:
                cv2.putText(frame, "Lost", (10, 30), font, 0.7, (0, 0, 255), 2)
                tracking = False

        # 未跟踪时显示所有目标
        if not tracking:
            for box, cls in zip(boxes, class_names):
                x, y, w, h = box
                cv2.rectangle(frame, (x, y), (x + w, y + h), (255, 0, 0), 2)
                cv2.putText(frame, cls, (x, y - 10), font, 0.5, (255, 0, 0), 2)

        # 显示提示
        hint = "Right click to stop" if tracking else "Left click to select"
        cv2.putText(frame, hint, (10, 30), font, 0.7, (0, 255, 255), 2)

        # 显示信息
        cv2.putText(frame, f"Tracker: {args.tracker}", (10, frame.shape[0] - 30),
                    font, 0.5, (255, 255, 0), 2)
        cv2.putText(frame, f"FPS: {int(cap.get(cv2.CAP_PROP_FPS))}",
                    (10, frame.shape[0] - 10), font, 0.5, (255, 255, 0), 2)

        # 显示窗口
        cv2.imshow(window_name, frame)

        # 退出
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

    cap.release()
    cv2.destroyAllWindows()


if __name__ == "__main__":
    main()
